package org.xx.armory.spring5.mvc.captchacode;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.GetMapping;
import org.xx.armory.commons.ForLogging;

import javax.servlet.http.HttpSession;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

import static javax.imageio.ImageIO.write;
import static org.springframework.http.HttpHeaders.CACHE_CONTROL;
import static org.springframework.http.HttpHeaders.EXPIRES;
import static org.springframework.http.HttpHeaders.PRAGMA;
import static org.springframework.http.MediaType.IMAGE_PNG;
import static org.xx.armory.spring5.security.AbstractArmoryUserPasswordAuthenticationProvider.CAPTCHA_CODE_KEY;

/**
 * 用于生成和校验验证码。
 */
@Controller
public class CaptchaCodeImageController {
    @ForLogging
    private final Logger logger = LoggerFactory.getLogger(CaptchaCodeImageController.class);

    @Autowired
    private CaptchaCodeImageProvider captchaCodeImageProvider;

    @GetMapping("/captcha-code-image")
    public ResponseEntity<ByteArrayResource> captchaCodeImage(
            HttpSession session
    )
            throws IOException {
        final var headers = new HttpHeaders();
        headers.add(CACHE_CONTROL, "no-cache, no-store, must-revalidate");
        headers.add(PRAGMA, "no-cache");
        headers.add(EXPIRES, "0");

        final var r = this.captchaCodeImageProvider.createCodeAndImage();

        session.setAttribute(CAPTCHA_CODE_KEY, r.getFirst());

        byte[] imageContent;
        try (final var output = new ByteArrayOutputStream()) {
            write(r.getSecond(), "PNG", output);
            output.flush();
            imageContent = output.toByteArray();
        }

        return ResponseEntity
                .ok()
                .headers(headers)
                .contentLength(imageContent.length)
                .contentType(IMAGE_PNG)
                .body(new ByteArrayResource(imageContent, "Captcha code image"));
    }
}
